"""
nnunet/network_architecture/neural_network.py
    - predict_3D
        - _internal_predict_3D_2Dconv_tiled
        - 此函数在2D中无效，其作用是将Z轴一个个提取出来，送入后续步骤计算
            - _internal_predict_2D_2Dconv_tiled
            - 此函数去除z轴后剩下(c, w, h)所以会自动添加一个维度(b, c, w, h)
            - 所以实际上2D如果不进行Crop和Resample，可以直接当作常规处理步骤执行
                - pad_nd_image
                    - res = np.pad(image, pad_list, mode, **kwargs)
                    - 此处对其patch_size, 如果输入和patch_size一致或者比这个大可以忽略
                - _compute_steps_for_sliding_window
                    - 计算整个目标时，需要的几个特定坐标位置
                - _internal_maybe_mirror_and_pred_2D
                    - 执行四次翻转，并送入模型进行推理，然后将结果融合
                    - inference_apply_nonlin
                        - 这句话等价于F.softmax(network(x), 1)
                    - 注意，当需要多次滑窗时，作者采用了高斯进行平滑处理
                        - result_torch[:, :] *= mult

"""
import cv2
import numpy
import torch
import torch.nn.functional as F
from Preprocess2D import *

model = torch.jit.load("v1_jit.pth")
model.eval()
sample_image = "../resource/raw_resize/1.jpg"

def get_sample_input():
    mat = cv2.imread(sample_image)
    mat = mat.transpose((2, 0, 1))
    mat = np.expand_dims(mat, 1)
    mat = mat.astype(np.float32)
    data, seg, _ = crop_to_nonzero(mat)
    out_data = data_normalize(data, seg)
    print(out_data.shape)
    return out_data

def simple_impl_infer_way():
    # 由于看不懂滑窗执行的过程，所以首先简化编写
    numpy_input = get_sample_input()
    numpy_input = numpy.transpose(numpy_input, (1, 0, 2, 3))

    # 执行多视角预测，并将结果进行softmax
    with torch.no_grad():
        x = torch.from_numpy(numpy_input)
        pred = F.softmax(model(x), 1)
        pred += torch.flip(F.softmax(model(torch.flip(x, (3, ))), 1), (3, ))
        pred += torch.flip(F.softmax(model(torch.flip(x, (2, ))), 1), (2, ))
        pred += torch.flip(F.softmax(model(torch.flip(x, (3, 2))), 1), (3, 2))
        pred /= 4
        pred = pred.argmax(1)

    pred = pred.cpu().numpy()[0]
    cv2.imwrite("out_pred.png", pred * 255)

if __name__ == '__main__':
    simple_impl_infer_way()